/**
 * @file Multigrid_Solver_2D.h
 * @brief A simple implementation for Multigrid Solver; 
 * It deals with 2-D Possion Equation on [0,1]*[0,1] with Dirichlet Boundary Condition;
 * @author Linhuo (1066117119@qq.com)
 * @version 1.0
 * @date 2021-06-07
 * 
 * @copyright Copyright (c) 2021  Linhuo
 * 
 */
#ifndef _PRO2_MULTIGRID_SOLVER_2D_H_
#define _PRO2_MULTIGRID_SOLVER_2D_H_
#include "Multigrid_Solver.h"
/**
 * @brief A simple implementation for Multigrid Solver; 
 * It deals with 2-D Possion Equation on [0,1]*[0,1] with Dirichlet Boundary Condition;
 */
class MG_2D : public Multigrid_Solver
{
public:
    void Initialize_u();
    void Initialize_BC(Function _BCF);
    void InitializeRhs(Function _F);
    VectorXd SmoothingforVC(VectorXd _u, VectorXd _f, double _h);
    VectorXd ResidualSolver(const VectorXd _u, const VectorXd _f, const double _h);
    VectorXd RestrictSolver(VectorXd _r);
    VectorXd IntpSolver(VectorXd _r);
    VectorXd PrepareFMG(VectorXd _f);
    int get_now_num_grid(VectorXd _f) { return (sqrt(_f.rows()) - 1); };
};

VectorXd MG_2D::PrepareFMG(VectorXd _r)
{
    VectorXd __u;
    int N = sqrt(_r.rows()) - 1;
    __u.setZero((N + 1) * (N + 1));
    for (int i = 0; i <= N; i++)
    {
        __u(i) = _BC(i, 0);
        __u((N + 1) * N + i) = _BC(i, 1);
        __u(i * (N + 1)) = _BC(i, 2);
        __u((i + 1) * (N + 1) - 1) = _BC(i, 3);
    };
    return __u;
};

void MG_2D::Initialize_u()
{
    int N = _num_of_grid;
    _u_origin.setZero((1 + N) * (1 + N));
    for (int i = 0; i <= N; i++)
    {
        _u_origin(i) = _BC(i, 0);
        _u_origin((N + 1) * N + i) = _BC(i, 1);
        _u_origin(i * (N + 1)) = _BC(i, 2);
        _u_origin((i + 1) * (N + 1) - 1) = _BC(i, 3);
    };
};

void MG_2D::Initialize_BC(Function _BCF)
{
    int num = _num_of_grid;
    _BC.setZero(num + 1, 4);
    double a[2] = {0, 0};
    ///The 0st col：Boundary y=0
    ///Yhe 1st col: Boundary y=1
    for (int j = 0; j <= 1; j++)
    {
        a[1] = j;
        for (int i = 0; i <= num; i++)
        {
            a[0] = 1.0 * i / num;
            _BC(i, j) = _BCF(a);
        };
    };
    ///The 2rd col：Boundary x=0
    ///Yhe 3rd col: Boundary x=1
    for (int j = 0; j <= 1; j++)
    {
        a[0] = j;
        for (int i = 0; i <= num; i++)
        {
            a[1] = 1.0 * i / num;
            _BC(i, j + 2) = _BCF(a);
        };
    };
};

void MG_2D::InitializeRhs(Function _F)
{
    int N = _num_of_grid;
    double a[2] = {0, 0};
    _f_origin.setZero((N + 1) * (N + 1));
    for (int j = 1; j < N; j++)
    {
        a[1] = 1.0 * j / N;
        for (int i = 1; i < N; i++)
        {
            a[0] = 1.0 * i / N;
            _f_origin(j * (N + 1) + i) = _F(a);
            // std::cout<<j*(N+1)+i<<": "<< a[0]<< a[1]<<"  "<<_F(a) <<std::endl;
        };
    };
    for (int i = 0; i <= N; i++)
    {
        _f_origin(i) = _BC(i, 0);
        _f_origin((N + 1) * N + i) = _BC(i, 1);
        _f_origin(i * (N + 1)) = _BC(i, 2);
        _f_origin((i + 1) * (N + 1) - 1) = _BC(i, 3);
    };
};

VectorXd MG_2D::SmoothingforVC(VectorXd _u, VectorXd _f, double _h)
{
    VectorXd _u_temp;
    int N = sqrt(_u.rows()) - 1;
    // std::cout << "We have N with " << N << std::endl;
    _u_temp.setZero((N + 1) * (N + 1));
    /// keep ghost cell the same in _u_temp.
    for (int i = 0; i <= N; i++)
    {
        _u_temp(i) = _u(i);
        _u_temp((N + 1) * N + i) = _u((N + 1) * N + i);
        _u_temp(i * (N + 1)) = _u(i * (N + 1));
        _u_temp((i + 1) * (N + 1) - 1) = _u((i + 1) * (N + 1) - 1);
    };

    double _t = 0;
    double omega = 2.0 / 3;
    switch (SmoothType)
    {
    case GaussSeidel:
        for (int j = 1; j <= N - 1; j++)
        {
            for (int i = 1; i <= N - 1; i++)
            {
                _u_temp(j * (N + 1) + i) = (-_h * _h * _f(j * (N + 1) + i) + _u_temp((j - 1) * (N + 1) + i) + _u((j + 1) * (N + 1) + i) + _u_temp(j * (N + 1) + i - 1) + _u(j * (N + 1) + i + 1)) / 4.0;
            };
        };
        break;
    case WeightedJacobi:
        for (int j = 1; j <= N - 1; j++)
        {
            for (int i = 1; i <= N - 1; i++)
            {
                _t = (-_h * _h * _f(j * (N + 1) + i) + _u((j - 1) * (N + 1) + i) + _u((j + 1) * (N + 1) + i) + _u(j * (N + 1) + i - 1) + _u(j * (N + 1) + i + 1)) / 4.0;
                _u_temp(j * (N + 1) + i) = (1.0 - omega) * _u(j * (N + 1) + i) + omega * _t;
            };
        };
        break;
    default:
        break;
    };
    return _u_temp;
};

VectorXd MG_2D::ResidualSolver(const VectorXd _u, const VectorXd _f, const double _h)
{
    VectorXd _u_temp;
    int N = sqrt(_u.rows()) - 1;
    _u_temp.setZero((N + 1) * (N + 1));
    for (int j = 1; j <= N - 1; j++)
    {
        for (int i = 1; i <= N - 1; i++)
        {
            _u_temp(j * (N + 1) + i) = _f(j * (N + 1) + i) - (_u((j - 1) * (N + 1) + i) + _u((j + 1) * (N + 1) + i) + _u(j * (N + 1) + i - 1) + _u(j * (N + 1) + i + 1) - 4.0 * _u(j * (N + 1) + i)) / (_h * _h);
        };
    };
    return _u_temp;
};

VectorXd MG_2D::RestrictSolver(VectorXd _r)
{
    VectorXd _u_temp;
    int N = (sqrt(_r.rows()) - 1) / 2;
    _u_temp.setZero((N + 1) * (N + 1));
    double _temp = 0;
    ///ghost cell
    for (int i = 0; i <= N; i++)
    {
        _u_temp(i) = _r(2 * i);
        _u_temp((N + 1) * N + i) = _r((2 * N + 1) * 2 * N + 2 * i);
        _u_temp(i * (N + 1)) = _r(2 * i * (2 * N + 1));
        _u_temp((i + 1) * (N + 1) - 1) = _r((2 * i + 1) * (2 * N + 1) - 1);
    };
    switch (RestricType)
    {
    case FullWeighting:
        for (int j = 1; j <= N - 1; j++)
        {
            for (int i = 1; i <= N - 1; i++)
            {
                _temp = (1.0 / 16.0) * (_r((2 * j - 1) * (2 * N + 1) + 2 * i - 1) + _r((2 * j - 1) * (2 * N + 1) + 2 * i + 1) + _r((2 * j + 1) * (2 * N + 1) + 2 * i - 1) + _r((2 * j + 1) * (2 * N + 1) + 2 * i + 1));
                _temp += (2.0 / 16.0) * (_r((2 * j - 1) * (2 * N + 1) + 2 * i) + _r((2 * j + 1) * (2 * N + 1) + 2 * i) + _r((2 * j) * (2 * N + 1) + 2 * i - 1) + _r((2 * j) * (2 * N + 1) + 2 * i + 1));
                _temp += (4.0 / 16.0) * (_r((2 * j) * (2 * N + 1) + 2 * i));
                _u_temp(j * (N + 1) + i) = _temp;
            };
        };
        break;
    case Injection:
        for (int j = 1; j <= N - 1; j++)
        {
            for (int i = 1; i <= N - 1; i++)
            {
                _u_temp(j * (N + 1) + i) = _r((2 * j) * (2 * N + 1) + 2 * i);
            };
        };
        break;
    default:
        break;
    };
    return _u_temp;
};

VectorXd MG_2D::IntpSolver(VectorXd _r)
{
    VectorXd _u_temp;
    int N = (sqrt(_r.rows()) - 1) * 2;
    _u_temp.setZero((N + 1) * (N + 1));
    if (_r(0) == _BC(0, 0) && _r(_r.rows() - 1) == _BC(_num_of_grid, 1))
    {
        for (int i = 0; i <= N; i++)
        {
            _u_temp(i) = _BC(i, 0);
            _u_temp((N + 1) * N + i) = _BC(i, 1);
            _u_temp(i * (N + 1)) = _BC(i, 2);
            _u_temp((i + 1) * (N + 1) - 1) = _BC(i, 3);
        };
    };
    switch (IntpType)
    {
    case Linear:
        for (int j = 1; j <= N - 1; j++)
        {
            for (int i = 1; i <= N - 1; i++)
            {
                if (i % 2 == 1 && j % 2 == 1)
                {
                    _u_temp(j * (N + 1) + i) = 1.0 / 4 * (_r((j - 1) / 2 * (N / 2 + 1) + (i - 1) / 2) + _r((j - 1) / 2 * (N / 2 + 1) + (i + 1) / 2) + _r((j + 1) / 2 * (N / 2 + 1) + (i - 1) / 2) + _r((j + 1) / 2 * (N / 2 + 1) + (i + 1) / 2));
                }
                else if (i % 2 == 0 && j % 2 == 0)
                {
                    _u_temp(j * (N + 1) + i) = _r(j / 2 * (N / 2 + 1) + i / 2);
                }
                else if (i % 2 == 0 && j % 2 == 1)
                {
                    _u_temp(j * (N + 1) + i) = 1.0 / 2 * (_r((j + 1) / 2 * (N / 2 + 1) + i / 2) + _r((j - 1) / 2 * (N / 2 + 1) + i / 2));
                }
                else if (i % 2 == 1 && j % 2 == 0)
                {
                    _u_temp(j * (N + 1) + i) = 1.0 / 2 * (_r((j) / 2 * (N / 2 + 1) + (i - 1) / 2) + _r((j) / 2 * (N / 2 + 1) + (i + 1) / 2));
                };
            };
        };
        break;
    case Quadratic:
        for (int j = 2; j <= N - 2; j++)
        {
            for (int i = 2; i <= N - 2; i++)
            {
                if (i % 2 == 0 && j % 2 == 0)
                {
                    _u_temp(j * (N + 1) + i) = _r(j / 2 * (N / 2 + 1) + i / 2);
                }
                else if (i % 2 == 0 && j % 2 == 1)
                {
                    _u_temp(j * (N + 1) + i) =
                        9.0 / 16 * (_r((j + 1) / 2 * (N / 2 + 1) + i / 2) + _r((j - 1) / 2 * (N / 2 + 1) + i / 2)) -
                        1.0 / 16 * (_r((j + 3) / 2 * (N / 2 + 1) + i / 2) + _r((j - 3) / 2 * (N / 2 + 1) + i / 2));
                }
                else if (j % 2 == 0 && i % 2 == 1)
                {
                    _u_temp(j * (N + 1) + i) =
                        9.0 / 16 * (_r((j) / 2 * (N / 2 + 1) + (i - 1) / 2) + _r((j) / 2 * (N / 2 + 1) + (i + 1) / 2)) -
                        1.0 / 16 * (_r((j) / 2 * (N / 2 + 1) + (i - 3) / 2) + _r((j) / 2 * (N / 2 + 1) + (i + 3) / 2));
                }
                else if (j % 2 == 1 && i % 2 == 1)
                {
                    _u_temp(j * (N + 1) + i) =
                        10.0 / 32 * (_r((j - 1) / 2 * (N / 2 + 1) + (i - 1) / 2) + _r((j - 1) / 2 * (N / 2 + 1) + (i + 1) / 2)) - 1.0 / 32 * (_r((j - 1) / 2 * (N / 2 + 1) + (i - 3) / 2) + _r((j - 1) / 2 * (N / 2 + 1) + (i + 3) / 2)) +
                        10.0 / 32 * (_r((j + 1) / 2 * (N / 2 + 1) + (i - 1) / 2) + _r((j + 1) / 2 * (N / 2 + 1) + (i + 1) / 2)) - 1.0 / 32 * (_r((j + 1) / 2 * (N / 2 + 1) + (i - 3) / 2) + _r((j + 1) / 2 * (N / 2 + 1) + (i + 3) / 2)) -
                        1.0 / 32 * (_r((j - 3) / 2 * (N / 2 + 1) + (i - 1) / 2) + _r((j - 3) / 2 * (N / 2 + 1) + (i + 1) / 2)) - 1.0 / 32 * (_r((j + 3) / 2 * (N / 2 + 1) + (i - 1) / 2) + _r((j + 3) / 2 * (N / 2 + 1) + (i + 1) / 2));
                    ;
                };
            };
        };

        for (int j = 2; j <= (N - 2); j += 2)
        {
            for (int i = 1; i <= (N - 1); i += (N - 2))
            {
                _u_temp(j * (N + 1) + i) = 1.0 / 2 * (_r((j) / 2 * (N / 2 + 1) + (i - 1) / 2) + _r((j) / 2 * (N / 2 + 1) + (i + 1) / 2));
            };
        };

        for (int j = 1; j <= (N - 1); j += (N - 2))
        {
            for (int i = 2; i <= (N - 2); i += 2)
            {
                _u_temp(j * (N + 1) + i) = 1.0 / 2 * (_r((j + 1) / 2 * (N / 2 + 1) + i / 2) + _r((j - 1) / 2 * (N / 2 + 1) + i / 2));
            };
        };
        for (int j = 1; j <= (N - 1); j += 2)
        {
            for (int i = 1; i <= (N - 1); i += (N - 2))
            {
                _u_temp(j * (N + 1) + i) = 1.0 / 4 * (_r((j - 1) / 2 * (N / 2 + 1) + (i - 1) / 2) + _r((j - 1) / 2 * (N / 2 + 1) + (i + 1) / 2) + _r((j + 1) / 2 * (N / 2 + 1) + (i - 1) / 2) + _r((j + 1) / 2 * (N / 2 + 1) + (i + 1) / 2));
            };
        };
        for (int j = 1; j <= (N - 1); j += (N - 2))
        {
            for (int i = 1; i <= (N - 1); i += 2)
            {
                _u_temp(j * (N + 1) + i) = 1.0 / 4 * (_r((j - 1) / 2 * (N / 2 + 1) + (i - 1) / 2) + _r((j - 1) / 2 * (N / 2 + 1) + (i + 1) / 2) + _r((j + 1) / 2 * (N / 2 + 1) + (i - 1) / 2) + _r((j + 1) / 2 * (N / 2 + 1) + (i + 1) / 2));
            };
        };

        break;
    default:
        break;
    };
    return _u_temp;
};

#else
//DO NOTHING
#endif